# prediction_gui.py
import sys
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                           QHBoxLayout, QComboBox, QPushButton, QLabel,
                           QFileDialog, QMessageBox)
from PyQt5.QtCore import Qt
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
import pandas as pd
from prediction import Prediction, StockDataset
from prediction import CNNModel  # 确保导入 CNNModel 类
import torch
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties
import matplotlib.font_manager as fm
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
import pandas as pd


# 在 predict 方法中，绘制图表之前添加以下代码
font_path = r"E:\simfang.ttf"  # 字体路径
prop = fm.FontProperties(fname=font_path)
# 设置字体
font = FontProperties(fname=r"E:\simfang.ttf", size=14)
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 设置字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

class PredictionGUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle('股票预测系统')
        self.setGeometry(100, 100, 1200, 800)

        # 初始化预测类
        self.prediction = Prediction(data_days=10, batch_size=200)
        self.dataset = StockDataset(data_days=10, remake_data=False)
        
        # 读取股票代码和名字对应表
        stock_names_df = pd.read_csv('data/stock_names.csv', header=None, names=['code', 'name'])

        # 确保股票代码是纯字符串，不带前缀和空格
        stock_names_df['code'] = stock_names_df['code'].astype(str).str.strip().str[-6:]

        # 转换为字典
        self.stock_names = stock_names_df.set_index('code')['name'].to_dict()

        # 创建主窗口部件和布局
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        layout = QVBoxLayout(main_widget)

        # 创建控制面板
        control_panel = QWidget()
        control_layout = QHBoxLayout(control_panel)
        
        # 添加模型选择下拉框
        self.model_combo = QComboBox()
        self.model_combo.addItems([ 'CNN', 'LSTM', 'RNN_RELU', 'ResNet18', 'ResNet34', 'ResNet50', 'ResNet101', 'ResNet152', 'DenseNet' ])
        control_layout.addWidget(QLabel('选择模型:'))
        control_layout.addWidget(self.model_combo)

        # 添加股票代码选择下拉框
        self.stock_combo = QComboBox()
        for code in self.dataset.stocks_codes[:20]:
            formatted_code = str(code).strip()[-6:]  # 取股票代码最后6位
            name = self.stock_names.get(formatted_code, '未知股票')
            self.stock_combo.addItem(f'{formatted_code} - {name}', formatted_code)

        control_layout.addWidget(QLabel('选择股票:'))
        control_layout.addWidget(self.stock_combo)

        # 添加预测按钮
        predict_btn = QPushButton('开始预测')
        predict_btn.clicked.connect(self.predict)
        control_layout.addWidget(predict_btn)

        # 添加控制面板到主布局
        layout.addWidget(control_panel)

        # 创建图表
        self.figure = plt.figure(figsize=(12, 6))
        self.canvas = FigureCanvas(self.figure)
        layout.addWidget(self.canvas)

    def predict(self):
        try:
            # 清除之前的图表
            self.figure.clear()

            # 获取选择的模型和股票代码
            model_name = self.model_combo.currentText()
            stock_code = self.stock_combo.currentData()  # 获取股票代码

            # 股票代码格式处理
            pure_code = str(stock_code).strip()[-6:]  # 提取6位数字代码
            
            # 尝试多种可能的文件命名格式
            possible_paths = [
                f'./data/stocks/{pure_code}.csv',
                f'./data/stocks/sh.{pure_code}.csv',
                f'./data/stocks/sz.{pure_code}.csv'
            ]
            
            df = None
            actual_path = None
            
            for path in possible_paths:
                if os.path.exists(path):
                    try:
                        df = pd.read_csv(path)
                        actual_path = path
                        print(f"成功读取文件: {path}")
                        break
                    except Exception as e:
                        print(f"尝试读取 {path} 失败: {str(e)}")
            
            if df is None:
                raise FileNotFoundError(f"找不到股票代码为 {pure_code} 的数据文件，已尝试路径: {possible_paths}")

            # 确定用于预测的股票代码格式
            prediction_stock_code = pure_code
            if 'sh.' in actual_path:
                prediction_stock_code = f'sh.{pure_code}'
            elif 'sz.' in actual_path:
                prediction_stock_code = f'sz.{pure_code}'
                
            # 获取股票名字
            stock_name = self.stock_names.get(pure_code, '未知股票')
            
            # 清理列名
            df.columns = df.columns.str.strip().str.lower()
            print("列名检查:", df.columns)

            # 确保必要的列存在
            if 'date' not in df.columns:
                raise ValueError("数据文件中缺少 'date' 列，请检查文件格式。")
            if 'close' not in df.columns:
                raise ValueError("数据文件中缺少 'close' 列，请检查文件格式。")

            # 提取交易日期和收盘价
            trading_dates = df['date'].values
            true_close = df['close'].values
            x_r = range(0, len(trading_dates))

            # 定义收盘价预测函数
            def close_p(x):
                if type(x) == int:
                    return x
                x = x[0, 1].item()
                return x if 0.2 > x > -0.2 else 0.0

            # 根据选择的模型进行预测
            predict_func = getattr(self.prediction, f'predict_{model_name.lower()}')
            predict_close = [true_close[j] * (1 + close_p(predict_func(prediction_stock_code, (0, j))))
                            for j in range(len(trading_dates))]

            # 绘制图表
            ax = self.figure.add_subplot(111)
            ax.plot(x_r, true_close, label='真实值')
            ax.plot(x_r, predict_close, label=f'{model_name}模型预测值')
            
            # 设置x轴标签
            x_ticks = list(x_r[::100])
            if x_r[-1] not in x_ticks:
                x_ticks.append(x_r[-1])
            x_labels = [trading_dates[i] for i in x_ticks]
            ax.set_xticks(x_ticks)
            ax.set_xticklabels(x_labels, rotation=45, fontproperties=prop)

            ax.set_ylabel('开盘价', fontproperties=prop)
            ax.legend(prop=prop)
            
            # 添加股票名字标签
            ax.set_title(f'{stock_name} ({pure_code}) - {model_name}模型预测', fontproperties=prop)
            
            # 更新画布
            self.canvas.draw()

        except Exception as e:
            import traceback
            traceback.print_exc()  # 打印完整错误堆栈
            QMessageBox.critical(self, '错误', f'预测过程中出现错误: {str(e)}')



if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = PredictionGUI()
    window.show()
    sys.exit(app.exec_())